from tkinter import Label
import warnings
import numpy as np
from scipy import sparse
from sklearn.semi_supervised import LabelPropagation, LabelSpreading
from sklearn.semi_supervised._label_propagation import BaseLabelPropagation
from sklearn.utils.multiclass import check_classification_targets
from sklearn.utils.extmath import safe_sparse_dot
from sklearn.exceptions import ConvergenceWarning
from scipy.sparse import csgraph

class MyBaseLabelPropagation(BaseLabelPropagation):
    def fit(self, X, y, graph_matrix=None):
        """Fit a semi-supervised label propagation model to X.

        The input samples (labeled and unlabeled) are provided by matrix X,
        and target labels are provided by matrix y. We conventionally apply the
        label -1 to unlabeled samples in matrix y in a semi-supervised
        classification.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        y : array-like of shape (n_samples,)
            Target class values with unlabeled points marked as -1.
            All unlabeled samples will be transductively assigned labels
            internally.

        graph_matrix: array-like of shape (n_samples, n_samples)

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        # X, y = self._validate_data(X, y) # 跳过类型检查
        self.X_ = X
        check_classification_targets(y)

        # actual graph construction (implementations should override this)
        # 如果 graph_matrix 直接提供，使用提供的 graph_matrix, 防止重复计算 graph_matrix
        if graph_matrix is None:
            graph_matrix = self._build_graph(self.X_)

        # label construction
        # construct a categorical distribution for classification only
        classes = np.unique(y)
        classes = classes[classes != -1]
        self.classes_ = classes

        n_samples, n_classes = len(y), len(classes)

        alpha = self.alpha
        if self._variant == "spreading" and (
            alpha is None or alpha <= 0.0 or alpha >= 1.0
        ):
            raise ValueError(
                "alpha=%s is invalid: it must be inside the open interval (0, 1)"
                % alpha
            )
        y = np.asarray(y)
        unlabeled = y == -1

        # initialize distributions
        self.label_distributions_ = np.zeros((n_samples, n_classes))
        for label in classes:
            self.label_distributions_[y == label, classes == label] = 1

        y_static = np.copy(self.label_distributions_)
        if self._variant == "propagation":
            # LabelPropagation
            y_static[unlabeled] = 0
        else:
            # LabelSpreading
            y_static *= 1 - alpha

        l_previous = np.zeros((self.X_.shape[0], n_classes))

        unlabeled = unlabeled[:, np.newaxis]
        if sparse.isspmatrix(graph_matrix):
            graph_matrix = graph_matrix.tocsr()

        for self.n_iter_ in range(self.max_iter):
            if np.abs(self.label_distributions_ - l_previous).sum() < self.tol:
                break

            l_previous = self.label_distributions_
            self.label_distributions_ = safe_sparse_dot(
                graph_matrix, self.label_distributions_
            )

            if self._variant == "propagation":
                normalizer = np.sum(self.label_distributions_, axis=1)[:, np.newaxis]
                normalizer[normalizer == 0] = 1
                self.label_distributions_ /= normalizer
                self.label_distributions_ = np.where(
                    unlabeled, self.label_distributions_, y_static
                )
            else:
                # clamp
                self.label_distributions_ = (
                    np.multiply(alpha, self.label_distributions_) + y_static
                )
        else:
            warnings.warn(
                "max_iter=%d was reached without convergence." % self.max_iter,
                category=ConvergenceWarning,
            )
            self.n_iter_ += 1

        normalizer = np.sum(self.label_distributions_, axis=1)[:, np.newaxis]
        normalizer[normalizer == 0] = 1
        self.label_distributions_ /= normalizer

        # set the transduction item
        transduction = self.classes_[np.argmax(self.label_distributions_, axis=1)]
        self.transduction_ = transduction.ravel()
        return self

class MyLabelPropagation(LabelPropagation, MyBaseLabelPropagation):
    """
    略微修改一下 fit 方法, 越过 X 的类型检查
    """
    def _build_graph(self, X):
        """Matrix representing a fully connected graph between each sample

        This basic implementation creates a non-stochastic affinity matrix, so
        class distributions will exceed 1 (normalization may be desired).
        """
        if self.kernel == "knn":
            self.nn_fit = None
        affinity_matrix = self._get_kernel(X)
        normalizer = affinity_matrix.sum(axis=0)
        if sparse.isspmatrix(affinity_matrix):
            affinity_matrix.data /= np.diag(np.array(normalizer))
        else:
            affinity_matrix /= normalizer[:, np.newaxis]
        return affinity_matrix

    def fit(self, X, y, graph_matrix=None):
        return MyBaseLabelPropagation.fit(self, X, y, graph_matrix)


class MyLabelSpreading(LabelSpreading, MyBaseLabelPropagation):
    """
    略微修改一下 fit 以及 graph build 方法，越过 X 的类型检查
    """
    def _build_graph(self, X):
        """Graph matrix for Label Spreading computes the graph laplacian"""
        # compute affinity matrix (or gram matrix)
        if self.kernel == "knn":
            self.nn_fit = None
        n_samples = X.shape[0]
        affinity_matrix = self._get_kernel(X)
        laplacian = csgraph.laplacian(affinity_matrix, normed=True)
        laplacian = -laplacian
        if sparse.isspmatrix(laplacian):
            diag_mask = laplacian.row == laplacian.col
            laplacian.data[diag_mask] = 0.0
        else:
            laplacian.flat[:: n_samples + 1] = 0.0  # set diag to 0.0
        return laplacian

if __name__ == "__main__":
    a = MyLabelPropagation()
    b = MyLabelSpreading()
    a.fit()
    b.fit()
    print()